import os
import numpy as np
import tensorflow as tf
from tensorflow.keras.preprocessing.image import load_img, img_to_array
import matplotlib.pyplot as plt
from PIL import ImageOps, Image

# 禁用GPU
os.environ['CUDA_VISIBLE_DEVICES'] = '-1'

def display_prediction_and_mask():
    """显示预测结果和分割掩码"""
    try:
        # 加载模型
        model = tf.keras.models.load_model('model.h5')
        print("模型加载成功")
        
        # 加载测试图像
        image_path = 'D:/project/dataset/segdata/images/Abyssinian_46.jpg'
        testimg = load_img(image_path, target_size=(160, 160))
        
        # 确保图像是RGB格式
        if testimg.mode != 'RGB':
            testimg = testimg.convert('RGB')
        
        # 预处理图像
        img_array = img_to_array(testimg)
        img_array = img_array / 255.0  # 归一化
        img_array = np.expand_dims(img_array, axis=0)  # 添加批次维度
        
        # 执行预测
        val_preds = model.predict(img_array)
        
        # 获取指定索引的预测结果
        pred_index = 0  # 使用索引0而不是10，因为我们的批次大小是1
        pred = val_preds[pred_index]
        mask = np.argmax(pred, axis=-1)
        
        # 创建彩色掩码以便更好地可视化
        colored_mask = create_colored_mask(mask)
        
        # 创建灰度掩码
        gray_mask = create_grayscale_mask(mask)
        
        # 显示结果
        fig, axes = plt.subplots(2, 2, figsize=(12, 12))
        
        # 显示原始图像
        axes[0, 0].imshow(testimg)
        axes[0, 0].set_title('原始图像')
        axes[0, 0].axis('off')
        
        # 显示预测结果的某个通道（例如第一个类别）
        axes[0, 1].imshow(pred[:, :, 0], cmap='viridis')
        axes[0, 1].set_title('预测结果（类别0）')
        axes[0, 1].axis('off')
        
        # 显示彩色分割掩码
        axes[1, 0].imshow(colored_mask)
        axes[1, 0].set_title('彩色分割掩码')
        axes[1, 0].axis('off')
        
        # 显示灰度分割掩码
        axes[1, 1].imshow(gray_mask, cmap='gray')
        axes[1, 1].set_title('灰度分割掩码')
        axes[1, 1].axis('off')
        
        plt.tight_layout()
        plt.show()
        
        # 保存结果
        gray_mask_img = Image.fromarray(gray_mask, mode='L')
        gray_mask_img.save('prediction_gray_mask.png')
        print("灰度分割掩码已保存为 prediction_gray_mask.png")
        
        colored_mask_img = Image.fromarray(colored_mask)
        colored_mask_img.save('prediction_colored_mask.png')
        print("彩色分割掩码已保存为 prediction_colored_mask.png")
        
        # 打印一些统计信息
        print(f"预测结果形状: {pred.shape}")
        print(f"分割掩码形状: {mask.shape}")
        print(f"分割掩码唯一值: {np.unique(mask)}")
        
    except Exception as e:
        print(f"显示预测结果时出错: {e}")

def create_colored_mask(mask):
    """为分割掩码创建彩色表示"""
    # 定义颜色映射
    colors = [
        [0, 0, 0],      # 类别0 - 黑色 (背景)
        [255, 0, 0],    # 类别1 - 红色
        [0, 255, 0],    # 类别2 - 绿色
        [0, 0, 255],    # 类别3 - 蓝色
    ]
    
    # 创建彩色掩码
    height, width = mask.shape
    colored_mask = np.zeros((height, width, 3), dtype=np.uint8)
    
    # 应用颜色到掩码
    for i in range(min(len(colors), np.max(mask) + 1)):
        colored_mask[mask == i] = colors[i]
    
    return colored_mask

def create_grayscale_mask(mask):
    """为分割掩码创建灰度表示"""
    # 定义灰度值映射
    gray_values = [0, 85, 170, 255]
    
    # 创建灰度掩码
    height, width = mask.shape
    gray_mask = np.zeros((height, width), dtype=np.uint8)
    
    # 应用灰度值到掩码
    for i in range(min(len(gray_values), np.max(mask) + 1)):
        gray_mask[mask == i] = gray_values[i]
    
    return gray_mask

if __name__ == '__main__':
    display_prediction_and_mask()